''' 计算费雪信息矩阵（2021.3.15王耀）'''

import torch
import numpy as np
import torch.nn.functional as F
import math

def loglikelihood(logots, labels):
    log_likelihood = torch.sum(F.binary_cross_entropy_with_logits(input=logots, target=labels))
    return log_likelihood

def get_FIM(log_likelihood, W):
    dW = torch.gradients(-log_likelihood, W)
    FIM = torch.matmul(torch.reshape(dW, (W.shape[0])), torch.reshape(dW, (W.shape[0])).t())
    return FIM


def IW(log_likelihood, W, bata, lamda, k):
    mid_ = bata / (2 * lamda * lamda)
    norm = torch.norm(W, 1)
    H = get_FIM(log_likelihood, W)
    I = torch.mul(torch.eye(H.shape[0]), mid_)

    conviance = torch.inverse(torch.add(H, I))

    trace = torch.trace(torch.mul(conviance), bata / 2)

    KLPQ = (norm + trace) * (1 / (2 * lamda * lamda)) - k + 0.5 * k * (1 / mid_) + \
                                                        0.5 * math.log(torch.det(conviance))
    return KLPQ